import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import nnls, minimize
from scipy.interpolate import interp1d
from scipy.linalg import svd
import warnings

warnings.filterwarnings('ignore')

# **************Parameter settings*************
NumWavelengths = 11  # 用于重建的波长数
NumWavelengthsInterp = NumWavelengths * 100  # 用于绘图的内插波长数
NumGaussianBasis = 500  # 高斯基数目
MinLambda = 760  # 最小的波长
MaxLambda = 770  # 最大的波长
FWHMset = 1  # 高斯基的全宽半最大值，用最大波长归一化并乘以1000，无单位

# 加载数据
ResponseFilename = np.loadtxt('data.txt')  # 响应度随不同Vg波长的变化(探测器)
ResponseCurves = ResponseFilename / np.max(ResponseFilename)  # 将输入响应曲线归一化

MeasuredSignals = np.loadtxt('data2.txt')

# **************Initialization of variables*************
NumRow = ResponseCurves.shape[0]  # Vg的数量(检测器)
NumCol = MeasuredSignals.shape[1] if MeasuredSignals.ndim > 1 else 1  # 从MeasuredSignals矩阵中处理多个数据

if MeasuredSignals.ndim == 1:
    MeasuredSignals = MeasuredSignals.reshape(-1, 1)

SimulatedSignals = np.zeros((NumRow, NumCol))  # 模拟信号的初始化
HiResReconstructedSpectrum = np.zeros((NumWavelengthsInterp, NumCol))  # 初始化重建谱，用于绘图

VecOfLambdas = np.linspace(MinLambda, MaxLambda, NumWavelengths) / MaxLambda  # 归一化波长
VecOfLambdasPlot = np.linspace(MinLambda, MaxLambda, NumWavelengthsInterp) / MaxLambda  # 归一化插值波长

HiResResponseCurves = np.zeros((NumRow, NumWavelengthsInterp))
for r in range(NumRow):
    f = interp1d(VecOfLambdas, ResponseCurves[r, :], kind='cubic', fill_value='extrapolate')
    HiResResponseCurves[r, :] = f(VecOfLambdasPlot)  # 样条插值

MeasuResidual = np.zeros(NumCol)  # 一种记录信号残差2范数的阵列
ReguTerm = np.zeros(NumCol)  # 记录正则化项的2范数的数组

# **************Reconstruction*************
GaussianCenter = np.linspace(MinLambda, MaxLambda, NumGaussianBasis) / MaxLambda  # 归一化波长中心
HiResGaussianBasis = np.zeros((NumWavelengthsInterp, NumGaussianBasis))  # 插值高斯基

GaussianSigma = FWHMset / 1000 / (2 * np.sqrt(2 * np.log(2)))  # 归一化的高斯基的标准差

# 生成所有的高斯基
for j in range(NumGaussianBasis):
    HiResGaussianBasis[:, j] = np.exp(-0.5 * ((VecOfLambdasPlot - GaussianCenter[j]) / GaussianSigma) ** 2)
    HiResGaussianBasis[:, j] = 1 / (GaussianSigma * np.sqrt(2 * np.pi)) * HiResGaussianBasis[:, j]


def FindOptGamma_GCV(c, A):
    """
    GCV method, see: Discrete inverse problems insight and algorithms, pp. 96.
    """
    U, sigma, Vh = svd(A, full_matrices=False)
    m = A.shape[0]

    def computeGCV(Gamma):
        fi = sigma ** 2 / (sigma ** 2 + Gamma ** 2)  # compute filter factor
        rho = 0
        for i in range(len(sigma)):
            beta = np.dot(U[:, i], c)
            rho += ((1 - fi[i]) * beta) ** 2
        return rho / (m - np.sum(fi)) ** 2

    # 使用优化方法找到最优的Gamma
    result = minimize(computeGCV, 1e-6, bounds=[(0, None)], method='L-BFGS-B')
    return result.x[0]


# 对每一列数据进行重建
for i in range(NumCol):
    # 将测量信号归一化
    MeasuredSignals_col = MeasuredSignals[:, i] / np.max(MeasuredSignals[:, i])

    # 计算权重矩阵
    WeightMatrix = np.dot(HiResResponseCurves, HiResGaussianBasis)

    # 使用GCV方法求出最优正则化参数
    OptimalGamma = FindOptGamma_GCV(MeasuredSignals_col, WeightMatrix)

    # 生成增广矩阵
    LaplacianMatrix = np.eye(NumGaussianBasis)  # 生成拉普拉斯矩阵，它是一个单位矩阵
    AugWeightMatrix = np.vstack([WeightMatrix, OptimalGamma ** 2 * LaplacianMatrix])
    AugMeasuredSignals = np.concatenate([MeasuredSignals_col, np.zeros(NumGaussianBasis)])

    # 最小二乘解，带非负约束
    GaussianCoefficients, resnorm = nnls(AugWeightMatrix, AugMeasuredSignals)

    # 重建的光谱
    HiResReconstructedSpectrum[:, i] = np.dot(HiResGaussianBasis, GaussianCoefficients)
    HiResReconstructedSpectrum[:, i] = HiResReconstructedSpectrum[:, i] / np.max(
        HiResReconstructedSpectrum[:, i])  # 归一化

    # 模拟信号
    SimulatedSignals[:, i] = np.dot(HiResResponseCurves, np.dot(HiResGaussianBasis, GaussianCoefficients))

    # 计算残差和正则化项
    MeasuResidual[i] = np.linalg.norm(SimulatedSignals[:, i] - MeasuredSignals_col) ** 2  # 2-范数的信号残差
    ReguTerm[i] = np.linalg.norm(np.dot(LaplacianMatrix, GaussianCoefficients)) ** 2  # 2-范数的正则化项

# 数据可视化
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
for i in range(NumCol):
    plt.plot(SimulatedSignals[:, i], '*', label=f'Signal {i + 1}')
plt.xlabel('V_g index')
plt.ylabel('a.u.')
plt.title('Simulated signals')
plt.legend()

plt.subplot(1, 2, 2)
for i in range(NumCol):
    plt.plot(VecOfLambdasPlot * MaxLambda, HiResReconstructedSpectrum[:, i], '*', label=f'Spectrum {i + 1}')
plt.title('Reconstructed spectrum')
plt.xlabel('λ (nm)')
plt.ylabel('a.u.')
plt.legend()

plt.tight_layout()
plt.show()

# 输出结果摘要
print("重建完成!")
print(f"信号残差范数: {MeasuResidual}")
print(f"正则化项范数: {ReguTerm}")

import numpy as np

# 假设HiResReconstructedSpectrum是重建的光谱数据矩阵，形状为 (NumWavelengthsInterp, NumCol)
# VecOfLambdasPlot * MaxLambda 是对应的波长数组

# 先组合波长和光谱数据
wavelengths = VecOfLambdasPlot * MaxLambda  # 波长，单位nm
spectra = HiResReconstructedSpectrum  # 重建光谱数据，归一化后

# 如果有多组光谱数据，按列保存
# 先构造一个二维数组，第一列是波长，后面是各组光谱
output_data = np.column_stack([wavelengths] + [spectra[:, i] for i in range(spectra.shape[1])])

# 保存为csv文件，文件名可自定义
np.savetxt('reconstructed.csv', output_data, delimiter=',',
           header=','.join(['Wavelength_nm'] + [f'Spectrum_{i+1}' for i in range(spectra.shape[1])]),
           comments='')

print("重建光谱数据已保存到 reconstructed_spectra.csv")
